#include "reflection_pad1d_tiling.h"
#include "kernel_operator.h"
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;
template <typename T>
class ReflectionPad1d {
public:
    __aicore__ inline ReflectionPad1d() {} 

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, ReflectionPad1dTilingData tiling) {
        this->left = tiling.padLeft;
        this->right = tiling.padRight;
        this->wSize = tiling.wSize;
        this->outWSize = this->wSize + this->left + this->right;
        this->alignWSize = tiling.alignWSize;
        this->alignOutWSize = tiling.alignOutWSize;
        
        this->blockIdx = GetBlockIdx();
        const uint32_t ncPerCore = tiling.ncPerCore;
        const uint32_t tailNC = tiling.tailNC;

        if (this->blockIdx < tailNC) {
            this->startNC = this->blockIdx * (ncPerCore + 1);
            this->endNC = this->startNC + (ncPerCore + 1);
        } 
        else {
            this->startNC = tailNC * (ncPerCore + 1) + (this->blockIdx - tailNC) * ncPerCore;
            this->endNC = this->startNC + ncPerCore;
        }

        // 绑定GM内存
        const uint32_t currentTaskCount = this->endNC - this->startNC;
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x) + this->startNC * this->wSize, currentTaskCount * this->wSize); 
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(y) + this->startNC * this->outWSize, currentTaskCount * this->outWSize);    

        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->alignWSize * sizeof(T)); 
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->outWSize * sizeof(T));
    }

    __aicore__ inline void Process() {
        for (uint32_t nc = startNC; nc < endNC; ++nc) {
            uint32_t localNcIdx = nc - startNC;
            uint32_t inGmOffset = localNcIdx * wSize;
            uint32_t outGmOffset = localNcIdx * outWSize;

            CopyIn(inGmOffset);
            Compute();
            CopyOut(outGmOffset);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t inGmOffset) {
        LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        DataCopyExtParams copyParams{1, static_cast<uint16_t>(this->wSize * sizeof(T)), 0, 0, 0};
        DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(this->alignWSize - this->wSize), 0};
        DataCopyPad(xLocal, xGm[inGmOffset], copyParams, padParams);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute() {
        LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();

        // 左填充
        for (int32_t i = 0; i < left; ++i) {
            int32_t refIdx = left - i;
            refIdx = refIdx < 0 ? 0 : refIdx;
            refIdx = Std::max(0, Std::min(refIdx, static_cast<int32_t>(wSize - 1)));
            T val = xLocal.GetValue(static_cast<uint32_t>(refIdx));
            yLocal.SetValue(static_cast<uint32_t>(i), val);
        }

        // 中间原始数据
        for (int32_t i = 0; i < wSize; ++i) {
            T val = xLocal.GetValue(static_cast<uint32_t>(i));
            yLocal.SetValue(static_cast<uint32_t>(left + i), val);
        }

        // 右填充
        for (int32_t i = 0; i < right; ++i) {
            int32_t refIdx = wSize - 2 - i;
            refIdx = refIdx < 0 ? 0 : refIdx;
            refIdx = Std::max(0, Std::min(refIdx, static_cast<int32_t>(wSize - 1)));
            T val = xLocal.GetValue(static_cast<uint32_t>(refIdx));
            yLocal.SetValue(static_cast<uint32_t>(left + wSize + i), val);
        }

        outQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(uint32_t outGmOffset) {
        LocalTensor<T> yLocal = outQueueY.DeQue<T>();
        DataCopyExtParams copyParams{1, static_cast<uint16_t>(this->outWSize * sizeof(T)), 0, 0, 0};
        DataCopyPad(yGm[outGmOffset], yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;   
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY; 
    GlobalTensor<T> xGm;   
    GlobalTensor<T> yGm;    

    uint32_t left = 0;      // 左填充量
    uint32_t right = 0;     // 右填充量
    uint32_t wSize = 0;    // 目标W维度大小
    uint32_t outWSize = 0;  // 输出W维度大小
    uint32_t alignWSize = 0;     // 目标维度32字节对齐后大小

    uint32_t startNC = 0;   // 起始N×C索引
    uint32_t endNC = 0;     // 结束N×C索引
    uint32_t blockIdx = 0;  // 当前核索引

};

extern "C" __global__ __aicore__ void reflection_pad1d(GM_ADDR x, GM_ADDR paddings, GM_ADDR y, ReflectionPad1dTilingData tiling) {
    switch(tiling.tilingKey) {
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP32):
        {
            ReflectionPad1d<float> op;
            op.Init(x, y, tiling);
            op.Process();
            break;
        }
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP16):
        {           
            ReflectionPad1d<half> op_fp16;
            op_fp16.Init(x, y, tiling);
            op_fp16.Process();
            break;
        }
        case static_cast<uint32_t>(TilingKey::TILING_KEY_BF16):
        {
            ReflectionPad1d<bfloat16_t> op_bf16;
            op_bf16.Init(x, y, tiling);
            op_bf16.Process();
            break;
        }
        default:
            // 不支持的数据类型，直接返回
            return;
    }
}
